#pragma once

#include <iostream>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <vector>
#include <future>
#include <functional>
#include <atomic>
#include <memory>

class ThreadPool
{
public:
	using Func = std::function<void()>;
	ThreadPool(int thread_num = 1) :_stop(false)
	{
		for (int i = 0; i < thread_num; i++)
		{
			_threads.emplace_back(&ThreadPool::entry, this);
		}
	}

	~ThreadPool()
	{
		Stop();
	}

	void Stop()
	{
		if (_stop) return;
		_stop = true;
		_cv.notify_all();
		for (auto& thread : _threads)
		{
			thread.join();
		}
	}

	template <class F, class ...Args>
	auto Push(F &&f, Args &&...args) -> std::future<decltype(f(args...))>
	{
		using return_type = decltype(f(args...));
		auto Func = std::bind(std::forward<F>(f), std::forward<Args>(args)...);
		auto task = std::make_shared<std::packaged_task<return_type()>>(Func);
		std::future<return_type> Fu = task->get_future();

		{
			std::unique_lock<std::mutex> lock(_mutex);
			_tasks.emplace_back([task]() { (*task)(); });
			_cv.notify_one();
		}

		return Fu;
	}

private:
	void entry()
	{
		while (!_stop)
		{
			std::vector<Func> tmp_tasks;
			{
				std::unique_lock<std::mutex> lock(_mutex);
				_cv.wait(lock, [this]() {
					return _stop || !_tasks.empty();
				});
				tmp_tasks.swap(_tasks);
			}
			for (auto& task : tmp_tasks)
			{
				task();
			}
		}
	}
private:
	std::atomic<bool> _stop;
	std::mutex _mutex;
	std::condition_variable _cv;
	std::vector<std::thread> _threads;
	std::vector<Func> _tasks;
};